import cv2
from cv2 import norm
import numpy as np

class HeatMask(object):
    def __init__(self, alpha, beta, gamma=0, heat_norm=False, userColor=cv2.COLORMAP_JET) -> None:
        """热力图掩膜

        Args:
            alpha (float): 原始图像透明度
            beta (float): 掩膜透明度
            gamma (float, optional): 叠加后图像亮度. Defaults to 0.
            heat_norm (bool, optional): 是否被归一化到0~255. Defaults to False.
            userColor (_type_, optional): _description_. Defaults to cv2.COLORMAP_JET.
        """
        self.userColor = userColor
        self.heat_norm = heat_norm
        self.beta = beta
        self.alpha = alpha
        self.gamma = gamma
    
    def __call__(self, org_img, heat_img, img_pth=None):
        """绘制热力图到图像上

        Args:
            org_img (np.ndarray): 输入图像BGR模式
            heat_img (np.ndarray): 输入热力图像1-D
            img_pth (str, optional): 输入保存图像路径png模式. Defaults to None.
        """
        if self.heat_norm:
            #norm_img = np.zeros(heat_img.shape)
            #cv2.normalize(heat_img, norm_img, 0, 255, cv2.NORM_MINMAX)
            hmax = heat_img.max()
            hmin = heat_img.min()
            hlen = hmax - hmin + 0.000001
            norm_img = 255 * (heat_img - hmin) / hlen
        else:
            norm_img = heat_img
        
        norm_img = norm_img.astype(np.uint8)
        
        heat = cv2.applyColorMap(norm_img, self.userColor)
        heat = cv2.cvtColor(heat, cv2.COLOR_BGR2RGB)

        heat_mask = cv2.addWeighted(org_img, self.alpha, heat, self.beta, self.gamma)

        if img_pth is not None:
            cv2.imwrite(img_pth, heat_mask)
        
        return heat_mask


if __name__ == '__main__':
    def gk(k, s):
        kx = cv2.getGaussianKernel(k, s)
        ky = cv2.getGaussianKernel(k, s)
        return np.multiply(kx, np.transpose(ky))
    heat_mask = gk(512, 512)
    heat_mask[heat_mask < 0] = 0
    print(heat_mask.shape)
    org_img = cv2.imread('lena_color.png', cv2.IMREAD_COLOR)

    heat = HeatMask(0.3, 0.7, heat_norm=True)
    heat(org_img, heat_mask, 'mask.png')

